package main.scala

/**
 * Algebraic rules.
 */
object SimplifyRewriter {

	import main.scala.AST._
	import org.kiama.rewriting.Rewriter._
	
	/**
	 * If expression @param e is a numeric, add it's
	 * value to element of according type in tuple @param s.
	 * The third element of the tuple are the elements left
	 * in the summation (i.e. not numerics).
	 * 
	 * Note: this reverse the order of the list when applied with fold.
	 */
	def sum_terms(s : Tuple3[Int,Double,List[Exp]])(e : Exp) = {
	  e match {
	    case NumI(v) => (s._1 + v, s._2, s._3)
	    case NumD(v) => (s._1, s._2 + v, s._3)
	    case _ => (s._1, s._2, e :: s._3)
	  }
	}
	
	/**
	 * If expression @param e is a numeric, multiply it's
	 * value to element of according type in tuple @param s.
	 */
	def mul_terms(s : Tuple3[Int,Double,List[Exp]])(e : Exp) = {
	  e match {
	    case NumI(v) => (s._1 * v, s._2, s._3)
	    case NumD(v) => (s._1, s._2 * v, s._3)
	    case Pow(NumI(v), NumI(-1)) => {
	      if (s._1 % v == 0) (s._1 / v, s._2, s._3)
	      else if (s._2 == 1.0) (s._1, s._2 / v, s._3) // TODO: here should return a fraction
	      else (s._1, s._2 / v, s._3)
	    }
	    case Pow(NumD(v), NumI(-1)) => (s._1, s._2 / v, s._3) // powers of zero are removed in "numerics"
	    case _ => (s._1, s._2, e :: s._3)
	  }
	}
	
	/**
	 * Rules for numerics that can be applied multiple times.
	 */
	val numerics = rule {
	  /**
	   * Move negation to the value, thus removing the Neg node.
	   */
	  case Neg(NumI(v)) => NumI(-v)
	  case Neg(NumD(v)) => NumD(-v)
	  
	  /**
	   * Remove neutral element of sum or product.
	   */
	  case Add(Nil) => NumI(0)
	  case Add(List(e)) => e
	  case Mul(Nil) => NumI(1)
	  case Mul(List(e)) => e
	  
	  /**
	   * Remove ^1 (integer) power.
	   */
	  case Pow(b, NumI(1)) => b
	  case Pow(_, NumI(0)) => NumI(1)
	  case Pow(NumI(0), _) => NumI(0)
	  case Pow(NumD(0), _) => NumI(0)
	}
	
	/**
	 * Rules for numerics that should be applied only one time.
	 */
	val numericsOnce = strategy {
	  /**
	   * Substitute sum of numerics with the result.
	   */
	  case Add(t) => {
	    val (i, d, rest) = t.foldLeft((0 : Int, 0.0 : Double, List[Exp]())) (sum_terms(_)(_))
	    (i, d, rest) match {
	      case (i, 0.0, Nil) => Some(NumI(i))
	      case (i, d  , Nil) => Some(NumD(d+i))
	      case (0, 0.0, List(e)) => Some(e)
	      // check length (i.e. some terms have been removed) to avoid rewriting the same term
	      case (i, 0.0, _) if (rest.length < t.length - 1) => Some(Add(NumI(i) :: rest))
	      case (i, d  , _) if (rest.length < t.length - 1) => Some(Add(NumI(d+i) :: rest))
	      case _ => None
	    }
	  }
	  /**
	   * Substitute product of numerics with result.
	   */
	  case Mul(t) => {
	    val (i, d, rest) = t.foldLeft((1 : Int, 1.0 : Double, List[Exp]())) (mul_terms(_)(_))
	    (i, d, rest) match {
	      case (0, _, _) | (_, 0.0, _) => Some(NumI(0))
	      case (i, 1.0, Nil) => Some(NumI(i))
	      case (i, d  , Nil) => Some(NumD(d*i))
	      case (1, 1.0, List(e)) => Some(e)
	      // check length (i.e. some terms have been removed) to avoid rewriting forms 2*a
	      case (i, 1.0, _) if (rest.length < t.length - 1) => Some(Mul(NumI(i) :: rest))
	      case (i, d  , _) if (rest.length < t.length - 1) => Some(Mul(NumD(d*i) :: rest))
	      case _ => None
	    }
	  }
	}

	/**
	 * Get coefficients of variables in @param list and sum them
	 */
	def getCount(list : List[Exp]): Double = {
	  var count = 0.0
	  for(x <- list) {
	   x match {
	     case NumI(i) => count = count + i
	     case NumD(d) => count = count + d
	     case Mul(xs) => count = count + getCount(xs)
	     case _ => 
	   } 
	  }
	  
	  if(count == 0.0)
	    count = 1.0
	    
	  count
	}

	/**
	 * Check if @param e contains variable
	 */
  def containsVariable(e: Exp): Boolean = {
    e match {
      case Var(_) => true
      case NumI(_) => false
      case NumD(_) => false
      case Neg(x) => containsVariable(x)
      case Pow(l,r) => containsVariable(l) || containsVariable(r)
      case Mod(l,r) => containsVariable(l) || containsVariable(r)
      case Add(xs) => {
        for (x <- xs) {
          if (containsVariable(x)) {
            return true
          }
        }
        false
      }
      case Mul(xs) => {
        for (x <- xs) {
          if (containsVariable(x)) {
            return true
          }
        }
        false
      }
    }
  }
	
  /**
   * Get variables in expression @param e
   */
  def getVariables(e: Exp): List[Var] = {
    var varList : List[Var] = List()
    e match {
      case Var(x) => List(Var(x))
      case NumI(_) => varList
      case NumD(_) => varList
      case Neg(x) => getVariables(x)
      case Pow(l,r) => (getVariables(l) ++ getVariables(r))
      case Mod(l,r) => (getVariables(l) ++ getVariables(r))
      case Add(xs) => {
        for (x <- xs) {
          varList = varList ++ getVariables(x)
        }
        varList
      }
      case Mul(xs) => {
        for (x <- xs) {
          varList = varList ++ getVariables(x)
        }
        varList
      }
    }
  }
  
  /**
   * Simplify and merge expressions @param ex1 and @param ex2
   */
  def mergeExpressions(ex1:Exp, ex2:Exp) : Option[Exp] = {
    (ex1,ex2) match {
      case (Var(x), Var(y)) => Some(Mul(List(NumI(2), Var(x))))
      case (Mul(xs), Mul(ys)) => Some(Mul((List(NumD(getCount(xs) + getCount(ys))) ++ getVariables(Mul(xs)))))
      case (Var(x), Mul(ys)) => Some(Mul(List(NumD(getCount(ys) + 1), Var(x))))
      case (Mul(xs), Var(y)) => Some(Mul(List(NumD(getCount(xs) + 1), Var(y))))
      case (Pow(l1,r1), Pow(l2,r2)) => Some(Pow(l1,Add(List(r1,r2))))
      case _ => None
    }
  }
  
  /**
   * Simplify and reduce expressions in @param xs
   */
  def reduceExpression(xs:List[Exp]) : Exp = {
    var exp = xs.head
    var tail = xs.tail
    for(t <- tail) {
      mergeExpressions(exp,t) match {
        case Some(x) => exp = x
        case None =>
      }
    }
   
    exp
  }

  /**
   * Check if @param xs and @param ys can be merged and simplified
   */
  def matchMulList(xs: List[Exp], ys: List[Exp]): Boolean = {
    for (x <- xs) {
      if (containsVariable(x)) {
        if (!containsExpression(x, ys)) {
          return false
        }
      }
    }

    true
  }

  /**
   * Check if expression @param e1 and @param e2 can be simplified
   */
  def reducibleExpressions(e1: Exp)(e2: Exp): Boolean = {
    if (compareExpressions(e1, e2)) {
      true
    } else {
      (e1, e2) match {
        case (Mul(xs), Mul(ys)) => matchMulList(xs, ys)
        case (Var(x), Mul(ys)) => containsExpression(Var(x), ys)
        case (Mul(ys), Var(x)) => containsExpression(Var(x), ys)
        case (Pow(l1, r1), Pow(l2, r2)) => {
          if(!containsVariable(l1) && !containsVariable(l2)) {
            compareExpressions(l1, l2)
          } else {
            false
          }
        }
        case _ => false
      }
    }
  }
	
  /**
   * Strategy to simplify variables
   */
  val variablesOnce = strategy {
    case Add(t) => {
      var (v, rest) = t.partition(containsVariable)
      if (v.length > 1) {
        while (v.length != 0) {
          var variable = v.head
          var list = v.filter(reducibleExpressions(variable))
          if (list.length > 1) {
            var newExp = reduceExpression(list)
            rest = newExp :: rest
          } else {
            rest = variable :: rest
          }

          v = v.filterNot(reducibleExpressions(variable))
          
        }

        if(compareExpressions(Mul(t), Mul(rest))) {
          None
        } else {
        	Some(Add(rest))
        }
      } else {
        None
      }
    }
    
    case Mul(t) => {
      var (v, rest) = t.partition((p: Exp) => p match { case Var(_) => true; case _ => false })
      if (v.length > 1) {
        while (v.length != 0) {
          var variable = v.head
          var list = v.filter(var1 => (variable, var1) match {
            case (Var(x), Var(y)) => 
              x == y
            case _ => false
          })
          if (list.length > 1) {
            rest = Pow(variable, NumI(list.length)) :: rest
          } else {
            rest = variable :: rest
          }

          v = v.filterNot(var1 => (variable, var1) match {
            case (Var(x), Var(y)) => 
              x == y
            case _ => false
          })
          
        }

        if(compareExpressions(Mul(t), Mul(rest))) {
          None
        } else {
        	Some(Mul(rest))
        }
      } else {
        None
      }
    }
  }

  /**
   * Check if @param xs and @param ys contains same expressions
   */
  def compareList(xs: List[Exp], ys: List[Exp]): Boolean = {
    if (xs.length == ys.length) {
      for (x <- xs) {
        if (!containsExpression(x, ys))
          return false
      }
      true
    } else {
      false
    }
  }

  /**
   * Check if @param exp is contained in @param xs
   */
  def containsExpression(exp:Exp, xs:List[Exp]) : Boolean = {
    for(x <- xs) {
      if(compareExpressions(exp, x))
        return true
    }
    
    false
  }
  
  /**
   * Compare expressions @param e1 and @param e2
   */
  def compareExpressions(e1:Exp, e2:Exp): Boolean = {
    (e1, e2) match {
      case (NumI(x), NumI(y)) => x == y
      case (NumD(x), NumD(y)) => x == y
      case (Var(x), Var(y)) => x == y
      case (Add(x), Add(y)) => (x.length == y.length) && compareList(x,y)
      case (Mul(x), Mul(y)) => (x.length == y.length) && compareList(x,y)
      case (Pow(r1,l1), Pow(r2,l2)) => compareExpressions(r1,r2) && compareExpressions(l1,l2)
      case (Mod(r1,l1), Mod(r2,l2)) => compareExpressions(r1,r2) && compareExpressions(l1,l2)
      case _ => false
    }
  }

  /**
   * Strategy to simplify modulo operator
   */
  val modulo = strategy {
    case Mod(NumI(x), NumI(y)) => {
      Some(NumI(x%y))
    }
    case Mod(x, y) => {
      if (compareExpressions(x, y)) {
        Some(NumI(0))
      } else {
        None
      }
    }
  }
  
  /**
   * Expand multiplications in @param xs and return sum of products
   */
  def expandMulList(xs: List[Exp]) : Add = {
    var listMul : List[Exp] = List()
    var listAdd : List[List[Exp]] = List()
    var list : List[List[Exp]] = List()

    for(x <- xs) {
      x match {
        case Add(ys) => {
          listAdd = ys :: listAdd
        }
        
        case x => listMul = x :: listMul
      }
    }
    
    list = listMul :: list
    
    for(xs <- listAdd) {
      for(l <- list) {
        for(x <- xs) {
          list = list.filterNot(var1 => var1 match {case y => l == y}) 
          list = (x::l) :: list
        }
      }
    }
    
    var listToAdd : List[Exp] = List()
    for(l <- list) {
      listToAdd = Mul(l) :: listToAdd
    }

    Add(listToAdd)
  }
  
  /**
   * Strategy to expand multiplication
   */
  val expandMul = strategy {
    case Mul(xs) => {
      val mulContainsAdd = xs.exists(x => x match {case Add(_) => true; case _ => false})
      if(mulContainsAdd) {
    	  Some(expandMulList(xs))
      } else {
        None
      }
    }
    case _ => None
  }
	
  val pow = rule {
    /* division (^{-1}) of numerics is handled in mul_terms */
    case Pow(NumI(x),NumI(y)) if y != -1 => NumI(scala.math.pow(x, y).toInt)
    
    case Pow(NumI(x), NumD(y)) => NumD(scala.math.pow(x,y))
    case Pow(NumD(x), NumI(y)) => NumD(scala.math.pow(x,y))
    case Pow(NumD(x), NumD(y)) => NumD(scala.math.pow(x,y))
  }
  
  /**
   * Strategy to put negation inside multiplication
   */
  val negStrategy = strategy {
    case Neg(Mul(x :: xs)) => {
      Some(Mul(Neg(x) :: xs))
    }
    case Neg(Var(x)) => {
      Some(Mul(List(Neg(NumI(1)), Var(x))))
    }
  }
  
  /**
   * Strategy to simplify Add inside Add, Mul inside Mul
   */
  val simplifyList = strategy {
    case Add(xs) => {
      var bool = false
      var list : List[Exp] = List()
      for(x <- xs) {
        x match {
          case Add(ys) => {
            bool = true
            list = ys ++ list
          }
          case y => list = y :: list
        }
      }
      if (bool)
        Some(Add(list))
      else
        None
    }
    
    case Mul(xs) => {
      var bool = false
      var list : List[Exp] = List()
      for(x <- xs) {
        x match {
          case Mul(ys) => {
            bool = true
            list = ys ++ list
          }
          case y => list = y :: list
        }
      }
      if (bool)
        Some(Mul(list))
      else
        None
    }
  }
	
	/**
	 * Entry point.
	 * Apply algebraic rules to simplify expression @param e.
	 */
	def start(e : Exp) : Exp = {
	  rewrite(reduce(
		  simplifyList
		  + expandMul
	      + numerics
	      + negStrategy
	      + numericsOnce
	      + variablesOnce
	      + modulo
	      + pow
	  ))(e)
	}
	
	/* This is taken from Kiama Rewriter.scala source. */
    def log[T] (s : => Strategy, msg : String, emitter : org.kiama.util.Emitter = new org.kiama.util.Emitter) : Strategy = 
    new Strategy {
        def apply (t1 : Term) = {
            val r = s (t1)
            r match {
                case Some (t2) =>
                    emitter.emitln (msg + t1 + " succeeded with " + t2)
                case None =>
//                    emitter.emitln (msg + t1 + " failed")
            }
            r
        }
    }

}
